package zuoban.memory;

import lombok.extern.log4j.Log4j;
import lombok.extern.slf4j.Slf4j;
import org.jetbrains.annotations.NotNull;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.util.CollectionUtils;

import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Slf4j
public class MyCustomMemory implements ChatMemory {

    private final Map<String,List<Message>> memoryMap = new ConcurrentHashMap<>();

    public MyCustomMemory() {
        log.info("MyCustomMemory init");
    }

    @Override
    public void add(@NotNull String conversationId,@NotNull List<Message> messages) {
        log.info("add memory: {},messages:{}", conversationId, messages);
        memoryMap.put(conversationId, messages);
    }

    @NotNull
    @Override
    public List<Message> get(@NotNull String conversationId) {
        List<Message> messages = memoryMap.get(conversationId);
        return CollectionUtils.isEmpty( messages) ? List.of() : messages;
    }

    @Override
    public void clear(@NotNull String conversationId) {
        memoryMap.remove(conversationId);
    }
}
